"""
    main
"""

from __future__ import absolute_import
from __future__ import print_function
from __future__ import division

import math
import os

import torch
import numpy as np
import random
import torch.optim as optim

from dynamic_pricing.common.cmd_args import cmd_args
from dynamic_pricing.common.consts import DEVICE
from dynamic_pricing.common.dataset import train_data, test_data
from dynamic_pricing.common.bipartite_graph import bg
from dynamic_pricing.common.recorder import cur_time
# from dynamic_pricing.common.recorder import dur_dist
from dynamic_pricing.model.deepcoevolve import DeepCoevolve
from tqdm import tqdm

from dynamic_pricing.user_cluster.classify_events import classify_events


def load_data(train_events, test_events):
    """

    :param train_events:
    :param test_events:
    :return:
    """
    train_data.load_events(train_events, 'train')
    test_data.load_events(test_events, 'test')

    for e_idx, cur_event in enumerate(test_data.ordered_events):
        cur_event.global_idx += train_data.num_events
        if cur_event.prev_user_event is None:
            continue
        train_events = train_data.user_event_lists[cur_event.user]
        if len(train_events) == 0:
            continue
        assert train_events[-1].t <= cur_event.t
        cur_event.prev_user_event = train_events[-1]
        cur_event.prev_user_event.next_user_event = cur_event

    print('# train:', train_data.num_events, '# test:', test_data.num_events)
    print('totally', cmd_args.num_users, 'users,', cmd_args.num_items, 'items')


def main_loop(HR, result_path, the_k, the_user_type_dic):
    """

    :param HR:
    :param result_path:
    :param the_k:
    :param the_user_type_dic:
    :return:
    """
    # 最佳测试集准确率
    best_result_test = 0
    # 最佳准确率的训练迭代次数
    the_best_epoch = -1

    bg.reset()
    for event in train_data.ordered_events:
        bg.add_event(event.user, event.item)
    # 重要的贡献
    model = DeepCoevolve(cmd_args.num_users,
                         cmd_args.num_items,
                         cmd_args.embed_dim,
                         k=the_k,
                         user_type_dic=the_user_type_dic,
                         score_func=cmd_args.score_func,
                         dt_type=cmd_args.dt_type,
                         max_norm=cmd_args.max_norm).to(
        DEVICE)
    optimizer = optim.Adam(model.parameters(), lr=cmd_args.learning_rate)  # 参数

    cur_time.reset(0)
    for e in train_data.ordered_events:
        cur_time.update_event(e.user, e.item, e.t)
    rc_dump = cur_time.dump()
    # test_mar = [0.0 for _ in range(cmd_args.num_epochs)]
    # test_mae = [0.0 for _ in range(cmd_args.num_epochs)]
    # test_rmse = [0.0 for _ in range(cmd_args.num_epochs)]

    # train_loss = [0.0 for _ in range(cmd_args.num_epochs)]
    # train_mae = [0.0 for _ in range(cmd_args.num_epochs)]
    # train_rmse = [0.0 for _ in range(cmd_args.num_epochs)]

    # 保存测试集准确率日志
    # f_accuracy = None
    # if result_path is not None:
    #     f_accuracy = open(result_path, 'a')
    #     f_accuracy.truncate(0)
    for epoch in range(cmd_args.num_epochs):
        # mae_temp = 0.0
        # rmse_temp = 0.0
        # loss_temp = 0.0

        cur_time.load_dump(*rc_dump)
        # mar, mae, rmse = model(train_data.ordered_events[-1].t,
        #                        test_data.ordered_events[:],
        #                        phase='test')

        # test_mar[epoch - 1], test_mae[epoch - 1], test_rmse[epoch - 1] = mar, mae, rmse
        #
        # print('测试集准确率：', round(mar * 100 / len(test_data.ordered_events[:]), 2), '%')
        # pbar = tqdm(range(cmd_args.iters_per_val))
        test_events = test_data.ordered_events[:3000]
        mar, loss_test, total = model(train_data.ordered_events[-1].t,
                                      test_events,
                                      phase='test',
                                      HR=HR)
        # test_mar[epoch - 1], test_mae[epoch - 1], test_rmse[epoch - 1] = mar

        # print('测试集准确率：', round(loss_test * 100 / len(test_events), 2), '%',
        #       '  ', loss_test, '/', len(test_events))
        p = []
        for i in range(len(loss_test)):
            if total[i] > 0:

                p.append([round(j / total[i], 2) for j in loss_test[i]])
            else:
                p.append([0])
        print('MAR:', mar, '测试集准确率HR：', p, '数量：', loss_test, total)

        # if loss_test * 100 / len(test_events) > best_result_test:
        #     best_result_test = loss_test * 100 / len(test_events)
        #     the_best_epoch = epoch - 1
        #
        # cur_epoch_result = 'epoch:' + str(epoch) + \
        #                    ' 测试集准确率：' + \
        #                    str(round(loss_test * 100 / len(test_events), 2)) + '%\n'
        # if result_path is not None:
        #     f_accuracy.write(cur_epoch_result)

        pbar = tqdm(range(cmd_args.iters_per_val))

        model.reset_user_events_list()
        begin_idx = 0
        cur_pbar = 0
        pbar_len = math.ceil(len(train_data.ordered_events) / cmd_args.bptt)
        for it in pbar:
            # cur_pos = np.random.randint(train_data.num_events - cmd_args.bptt)
            cur_pos = cur_pbar * pbar_len
            cur_pbar += 1
            T_begin = 0
            if cur_pos:
                T_begin = train_data.ordered_events[cur_pos - 1].t
            cur_end = min(len(train_data.ordered_events), cur_pos + cmd_args.bptt)
            event_mini_batch = train_data.ordered_events[cur_pos:cur_end]

            optimizer.zero_grad()
            # loss, mae, rmse = model(T_begin, event_mini_batch, phase='train')
            loss_test = model(T_begin, event_mini_batch, phase='train', HR=HR)
            pbar.set_description('epoch: %.2f, loss: %.2f' % (
                epoch + (it + 1) / len(pbar), loss_test.item()))
            # pbar.set_description('epoch: %.2f, loss: %.2f, mae: %.2f, rmse: %.2f' % (
            #     epoch + (it + 1) / len(pbar), loss.item(), mae, rmse))

            # loss_temp, mae_temp, rmse_temp = loss.item(), mae, rmse

            loss_test.backward()  # 参数梯度的回传

            if cmd_args.grad_clip > 0:  # 防止某些参数过大
                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=cmd_args.grad_clip)

            optimizer.step()  # 更新参数
            model.normalize()

        # train_loss[epoch - 1], train_mae[epoch - 1], train_rmse[epoch - 1] = loss_temp, mae_temp, rmse_temp
        # dur_dist.print_dist()
    # if result_path is not None:
    #     f_accuracy.close()
    return best_result_test, the_best_epoch

    # path_test = r'C:\Users\lvsk\Desktop\test_coevolve.txt'
    # path_train = r'C:\Users\lvsk\Desktop\train_coevolve.txt'
    # f_test = open(path_test, 'a')
    # f_train = open(path_train, 'a')
    # for i in range(cmd_args.num_epochs):
    #     train = str(test_mar[i]) + ' ' + str(test_mae[i]) + ' ' + str(test_rmse[i]) + '\n'
    #     test = str(train_loss[i]) + ' ' + str(train_mae[i]) + ' ' + str(train_rmse[i]) + '\n'
    #     f_train.write(train)
    #     f_test.write(test)


if __name__ == '__main__':
    random.seed(cmd_args.seed)
    np.random.seed(cmd_args.seed)
    torch.manual_seed(cmd_args.seed)

    # user聚类数
    for k in range(6, 7):
        # path1 = r'./best_result_file'
        # if not os.path.exists(path1):
        #     os.mkdir(path1)
        # best_result_file_path = r'./best_result_file/' + str(k) + '聚类.txt'
        # f = open(best_result_file_path, 'a')
        # f.truncate(0)
        events_class, user_type_dic = classify_events(k=k)
        # events_class = [events_class]
        # 对不同的类别的user相关事件分别进行训练与预测
        for events_idx in range(len(events_class)):
            events = events_class[events_idx]
            print('聚类数目为:', k, ' 当前训练第', events_idx, '类user,共', len(events_class), '类')
            if len(events) > 100:
                train_index = int(len(events) * 0.9)
                train_events_list = events[:train_index]
                test_events_list = events[train_index:]

                load_data(train_events_list, test_events_list)
                # path2 = r'./k@' + str(k)
                # if not os.path.exists(path2):
                #     os.mkdir(path2)
                # result_file_path = r'./k@' + str(k) + r'/' + str(events_idx) + '.txt'
                # best_test_result, best_epoch = main_loop(HR=5, result_path=result_file_path)
                best_test_result, best_epoch = main_loop(HR=5, result_path=None,
                                                         the_k=k, the_user_type_dic=user_type_dic)
                result = '第' + str(events_idx) + '类user最佳测试集结果为:  '
                result += '第' + str(best_epoch) + '迭代  '
                result += '准确率为' + str(best_test_result) + '%'
            else:
                result = '第' + str(events_idx) + '类user没有交互数据'
            print(result)

            # f.write(result + '\n')
        # f.close()
